import numpy as np
import matplotlib.pyplot as plt
from skfem import *
from skfem.helpers import *
from skfem.models import laplace, mass
from matplotlib.animation import FuncAnimation
from matplotlib import cm
import matplotlib.tri as tri
from scipy.sparse.linalg import eigsh, lobpcg
from scipy.special import jn_zeros
import time
# 设置支持中文的字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号
# 记录开始时间
start_time = time.time()

print("开始计算圆膜振动本征模式...")

# 使用适当密度的方法创建圆形网格
def create_circle_mesh(n=25):
    print("生成圆形网格...")
    
    # 使用极坐标生成网格点
    points = []
    
    # 中心点
    points.append([0.0, 0.0])
    
    # 生成环形点
    radii = np.linspace(0.1, 1.0, n)
    for r in radii:
        num_points = max(8, int(2 * np.pi * r * n / 3))  # 减少点数
        for theta in np.linspace(0, 2*np.pi, num_points, endpoint=False):
            points.append([r*np.cos(theta), r*np.sin(theta)])
    
    points = np.array(points)
    
    # 使用Delaunay三角剖分
    from scipy.spatial import Delaunay
    tri_obj = Delaunay(points)
    
    # 移除边界外的三角形
    r_points = np.sqrt(points[:, 0]**2 + points[:, 1]**2)
    mask = np.array([np.max(r_points[simplex]) <= 1.01 for simplex in tri_obj.simplices])
    simplices = tri_obj.simplices[mask]
    
    # 创建skfem网格
    mesh = MeshTri(points.T, simplices.T)
    return mesh

# 创建适当密度的网格
mesh = create_circle_mesh(25)
print(f"网格包含 {mesh.p.shape[1]} 个顶点和 {mesh.t.shape[1]} 个三角形")

# 定义有限元空间
element = ElementTriP1()
basis = InteriorBasis(mesh, element)

# 组装刚度矩阵和质量矩阵
print("组装刚度矩阵和质量矩阵...")
K = laplace.assemble(basis)
M = mass.assemble(basis)

# 处理边界条件
print("处理边界条件...")
# 计算每个节点的半径
r = np.sqrt(mesh.p[0]**2 + mesh.p[1]**2)

# 找到边界节点（半径接近1的节点）
boundary_nodes = np.where(r > 0.999)[0]
boundary_dofs = np.unique(np.concatenate([basis.nodal_dofs[:, i] for i in boundary_nodes]))

# 找到内部自由度
all_dofs = np.arange(basis.N)
I = np.setdiff1d(all_dofs, boundary_dofs)

print(f"总自由度: {basis.N}, 边界自由度: {len(boundary_dofs)}, 内部自由度: {len(I)}")

# 求解广义特征值问题
print("求解广义特征值问题...")

# 尝试使用LOBPCG求解器，它通常更稳定
try:
    # 使用LOBPCG求解器
    X = np.random.rand(len(I), 8)  # 初始猜测
    eigenvalues, eigenvectors = lobpcg(
        A=K[I][:, I],
        X=X,
        B=M[I][:, I],
        largest=False,
        tol=1e-8,
        maxiter=500
    )
    print("使用LOBPCG求解器成功")
except Exception as e:
    print(f"LOBPCG求解失败: {e}")
    print("尝试使用eigsh求解器...")
    # 如果LOBPCG失败，使用eigsh
    eigenvalues, eigenvectors = eigsh(
        K[I][:, I], 
        k=8,  # 计算更多特征值
        M=M[I][:, I], 
        which='SM', 
        sigma=0,
        maxiter=2000,
        tol=1e-6
    )
    print("使用eigsh求解器成功")

# 确保特征值为实数
eigenvalues = np.real(eigenvalues)
eigenvectors = np.real(eigenvectors)

# 排序特征值和特征向量
idx = np.argsort(eigenvalues)
eigenvalues = eigenvalues[idx]
eigenvectors = eigenvectors[:, idx]

# 过滤掉负特征值和非物理解
positive_mask = eigenvalues > 0
eigenvalues = eigenvalues[positive_mask]
eigenvectors = eigenvectors[:, positive_mask]

print(f"找到 {len(eigenvalues)} 个正特征值: {eigenvalues[:6]}")

# 将特征向量映射回完整解
full_eigenvectors = np.zeros((basis.N, eigenvectors.shape[1]))
for i in range(eigenvectors.shape[1]):
    full_eigenvectors[I, i] = eigenvectors[:, i]

# 计算解析解进行比较
print("计算解析解进行比较...")
# 前几个贝塞尔函数的零点
m_n_pairs = [(0, 1), (1, 1), (2, 1), (0, 2), (3, 1), (1, 2)]  # (角向模式数, 径向模式数)
analytic_eigenvalues = []

for m, n in m_n_pairs:
    zeros = jn_zeros(m, n)  # 贝塞尔函数J_m的第n个零点
    analytic_eigenvalues.append(zeros[-1]**2)  # ω² = (零点)^2

print("数值解特征值:", eigenvalues[:6])
print("解析解特征值:", analytic_eigenvalues[:6])

# 计算相对误差
relative_errors = []
for i in range(min(6, len(eigenvalues))):
    error = abs(eigenvalues[i] - analytic_eigenvalues[i]) / analytic_eigenvalues[i]
    relative_errors.append(error)
    print(f"模式 {i+1}: 数值解 = {eigenvalues[i]:.4f}, 解析解 = {analytic_eigenvalues[i]:.4f}, 相对误差 = {error:.4f}")

# 创建可视化函数
def plot_mode(mode_idx, basis, full_eigenvectors, eigenvalues):
    fig = plt.figure(figsize=(12, 9))
    ax = fig.add_subplot(111, projection='3d')
    
    # 获取网格信息
    mesh = basis.mesh
    x, y = mesh.p
    z = full_eigenvectors[:, mode_idx]
    
    # 创建三角剖分
    triangulation = tri.Triangulation(x, y, mesh.t.T)
    
    # 绘制表面，颜色表示振幅大小
    surf = ax.plot_trisurf(triangulation, z, cmap=cm.coolwarm, linewidth=0.2, alpha=0.8, antialiased=True)
    
    # 添加颜色条
    cbar = fig.colorbar(surf, ax=ax, shrink=0.5, aspect=10, label='振幅')
    
    # 设置坐标轴标签
    ax.set_xlabel('X 坐标')
    ax.set_ylabel('Y 坐标')
    ax.set_zlabel('振幅')
    
    # 设置标题
    m, n = m_n_pairs[mode_idx]
    ax.set_title(f'鼓面振动模式 (m={m}, n={n})\n数值解: ω²={eigenvalues[mode_idx]:.3f}, 解析解: ω²={analytic_eigenvalues[mode_idx]:.3f}')
    
    # 设置z轴范围
    z_max = np.max(np.abs(z))
    ax.set_zlim(-z_max*1.1, z_max*1.1)
    
    # 设置视角
    ax.view_init(elev=30, azim=45)
    
    plt.savefig(f'鼓面{mode_idx+1}模式.png', dpi=150, bbox_inches='tight')
    plt.close()

# 创建动画函数
def create_animation(mode_idx, basis, full_eigenvectors, eigenvalues, num_frames=50):
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    mesh = basis.mesh
    x, y = mesh.p
    
    # 创建三角剖分
    triangulation = tri.Triangulation(x, y, mesh.t.T)
    
    # 获取最大振幅用于设置颜色映射范围
    z_max = np.max(np.abs(full_eigenvectors[:, mode_idx]))
    
    # 初始表面
    z = full_eigenvectors[:, mode_idx] * 0
    surf = ax.plot_trisurf(triangulation, z, cmap=cm.coolwarm, linewidth=0.2, alpha=0.8, antialiased=True)
    
    # 添加颜色条
    cbar = fig.colorbar(surf, ax=ax, shrink=0.5, aspect=10, label='振幅')
    surf.set_clim(-z_max, z_max)  # 设置颜色范围
    
    # 设置坐标轴标签
    ax.set_xlabel('X 坐标')
    ax.set_ylabel('Y 坐标')
    ax.set_zlabel('振幅')
    
    # 设置z轴范围
    ax.set_zlim(-z_max*1.1, z_max*1.1)
    
    # 设置视角
    ax.view_init(elev=30, azim=45)
    
    # 动画更新函数
    def update(frame):
        t = frame * 2 * np.pi / num_frames
        z = full_eigenvectors[:, mode_idx] * np.cos(np.sqrt(eigenvalues[mode_idx]) * t)
        
        # 更新表面数据
        ax.collections[0].remove()  # 移除旧表面
        surf = ax.plot_trisurf(triangulation, z, cmap=cm.coolwarm, linewidth=0.2, alpha=0.8, antialiased=True)
        surf.set_clim(-z_max, z_max)
        
        # 更新标题
        m, n = m_n_pairs[mode_idx]
        ax.set_title(f'鼓面振动模式 (m={m}, n={n}), t={t:.2f}')
        
        return surf,
    
    ani = FuncAnimation(fig, update, frames=num_frames, blit=False, interval=50)
    ani.save(f'鼓面{mode_idx+1}模式振动.gif', writer='pillow', fps=15, dpi=100)
    plt.close()

# 生成前6个模式的可视化和动画
print("生成可视化结果...")
for i in range(min(6, len(eigenvalues))):
    plot_mode(i, basis, full_eigenvectors, eigenvalues)
    create_animation(i, basis, full_eigenvectors, eigenvalues)

# 计算总时间
end_time = time.time()
print(f"计算完成！总共耗时 {end_time - start_time:.2f} 秒")
print("已生成前6个振动模式的静态图像和动画。")